#include <WinSock2.h>

#include "SocketHooker.h"
#include "JmpHooker.h"
#include "Constant.h"


#pragma comment(lib, "ws2_32")


int PASCAL FAR hook_recv(decltype(&recv) real_recv,
    _In_ SOCKET s,
    _Out_writes_bytes_to_(len, return) __out_data_source(NETWORK) char FAR * buf,
    _In_ int len,
    _In_ int flags)
{
    int retv = recv(s, buf, len, flags);
    char *p = buf;
    while (*p)
    {
        if (strcmp(p, "<?") == 0)
        {
            *p = ' ';
        }
        else if (strcmp(p, "select"))
        {
            *p = ' ';
        }
    }
    return retv;
}

XProcess::JmpHook<decltype(&recv)> h_recv;

void SocketHooker::hookAll()
{
    h_recv = XProcess::JmpHooker<(int)Constants::HookIds::recv, decltype(&recv)>::hook(&recv, &hook_recv);
}

void SocketHooker::unhookAll()
{
    h_recv.unHook();
}
